# 另一种广泛使用的模型复用方法是模型微调（fine-tuning），与特征提取互为补充。对于用
# 于特征提取的冻结的模型基，微调是指将其顶部的几层“解冻”，并将这解冻的几层和新增加的
# 部分（本例中是全连接分类器）联合训练（见图 5-19）。之所以叫作微调，是因为它只是略微调
# 整了所复用模型中更加抽象的表示，以便让这些表示与手头的问题更加相关。
# 微调网络的步骤如下。
# (1) 在已经训练好的基网络（base network）上添加自定义网络。
# (2) 冻结基网络。
# (3) 训练所添加的部分。
# (4) 解冻基网络的一些层。
# (5) 联合训练解冻的这些层和添加的部分。


# ---------------01 将 VGG16 卷积基实例化---------------
from keras.applications import VGG16

conv_base = VGG16(weights='imagenet', include_top=False, input_shape=(150, 150, 3))
# conv_base.summary()
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #
# =================================================================
# input_1 (InputLayer)         (None, 150, 150, 3)       0
# _________________________________________________________________
# block1_conv1 (Conv2D)        (None, 150, 150, 64)      1792
# _________________________________________________________________
# block1_conv2 (Conv2D)        (None, 150, 150, 64)      36928
# _________________________________________________________________
# block1_pool (MaxPooling2D)   (None, 75, 75, 64)        0
# _________________________________________________________________
# block2_conv1 (Conv2D)        (None, 75, 75, 128)       73856
# _________________________________________________________________
# block2_conv2 (Conv2D)        (None, 75, 75, 128)       147584
# _________________________________________________________________
# block2_pool (MaxPooling2D)   (None, 37, 37, 128)       0
# _________________________________________________________________
# block3_conv1 (Conv2D)        (None, 37, 37, 256)       295168
# _________________________________________________________________
# block3_conv2 (Conv2D)        (None, 37, 37, 256)       590080
# _________________________________________________________________
# block3_conv3 (Conv2D)        (None, 37, 37, 256)       590080
# _________________________________________________________________
# block3_pool (MaxPooling2D)   (None, 18, 18, 256)       0
# _________________________________________________________________
# block4_conv1 (Conv2D)        (None, 18, 18, 512)       1180160
# _________________________________________________________________
# block4_conv2 (Conv2D)        (None, 18, 18, 512)       2359808
# _________________________________________________________________
# block4_conv3 (Conv2D)        (None, 18, 18, 512)       2359808
# _________________________________________________________________
# block4_pool (MaxPooling2D)   (None, 9, 9, 512)         0
# _________________________________________________________________
# block5_conv1 (Conv2D)        (None, 9, 9, 512)         2359808
# _________________________________________________________________
# block5_conv2 (Conv2D)        (None, 9, 9, 512)         2359808
# _________________________________________________________________
# block5_conv3 (Conv2D)        (None, 9, 9, 512)         2359808
# _________________________________________________________________
# block5_pool (MaxPooling2D)   (None, 4, 4, 512)         0
# =================================================================
# Total params: 14,714,688
# Trainable params: 14,714,688
# Non-trainable params: 0
# _________________________________________________________________

# -----------------------------------自定义网络-----------------------------
from keras import models
from keras import layers

model = models.Sequential()
model.add(conv_base)
model.add(layers.Flatten())
model.add(layers.Dense(256, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))

# ---------------准备数据---------------
import os
from keras.preprocessing.image import ImageDataGenerator

base_dir = '/home/python-test/py36-keras-demo01/cats_and_dogs_small'
train_dir = os.path.join(base_dir, 'train')
validation_dir = os.path.join(base_dir, 'validation')
test_dir = os.path.join(base_dir, 'test')

train_datagen = ImageDataGenerator(
    rescale=1. / 255,
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest')
test_datagen = ImageDataGenerator(rescale=1. / 255)  # 注意，不能增强验证数据
train_generator = train_datagen.flow_from_directory(
    train_dir,  # 目标目录
    target_size=(150, 150),  # 将所有图像的大小调整为 150×150
    batch_size=20,
    class_mode='binary'  # 因为使用了binary_crossentropy损失，所以需要用二进制标签
)
validation_generator = test_datagen.flow_from_directory(
    validation_dir,
    target_size=(150, 150),
    batch_size=20,
    class_mode='binary'
)

# -----------------------------------冻结直到某一层的所有层-----------------------------
conv_base.trainable = True
set_trainable = False
for layer in conv_base.layers:
    if layer.name == 'block5_conv1':
        set_trainable = True
    if set_trainable:
        layer.trainable = True
    else:
        layer.trainable = False

# -----------------------------------微调模型-----------------------------------
from keras import optimizers

model.compile(loss='binary_crossentropy',
              optimizer=optimizers.RMSprop(lr=1e-5),
              metrics=['acc'])

# -----------------------------------保存最好的模型-----------------------------------
from keras.callbacks import ModelCheckpoint

filepath = "07_dog_vs_cat_vgg16_微调模型-{epoch:02d}-{val_acc:.4f}.hdf5"
# ModelCheckpoint保存模型
checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                             mode='max')
# -----------------------------------使用一个 TensorBoard 回调函数来训练模型-----------------------------------
from keras.callbacks import TensorBoard

tensorboard_log_dir = './tmp/log'

tensorboard = TensorBoard(
    log_dir=tensorboard_log_dir,

)
callbacks_list = [checkpoint,tensorboard]

# -----------------------------------开始训练-----------------------------------
history = model.fit_generator(
    train_generator,
    steps_per_epoch=100,
    epochs=100,
    validation_data=validation_generator,
    validation_steps=50,
    callbacks=callbacks_list)

# ---------------绘制结果---------------
import matplotlib.pyplot as plt

acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(acc) + 1)
plt.plot(epochs, acc, 'bo', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()
plt.figure()
plt.plot(epochs, loss, 'bo', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show(block=False)
plt.show()
